#pragma once 

//implement by red-black tree
template<typename ValueType>
class TSet
{

    class RedBlackTreeNode
    {
    public:
        RedBlackTreeNode(RedBlackTreeNode* InLeft,
            RedBlackTreeNode* InRight,
            const ValueType& InValue,
           // uint64_t InHash,
            bool InIsRed)
            : Left(InLeft)
            , Right(InRight)
            , bIsRed(InIsRed)
            , Value(InValue)
            //, Hash(InHash)
        {

        }

        RedBlackTreeNode* Left = nullptr;
        RedBlackTreeNode* Right = nullptr;
        uint64_t Size = 1;   // Left size+ right size +1. for leaf, size is 1
        //uint64_t Hash = 0;
        bool bIsRed = true;
        ValueType Value;

        ~RedBlackTreeNode()
        {
            if (Left)
            {
                delete Left;
            }

            if (Right)
            {
                delete Right;
            }
        }

    };

public:

    void Add(const ValueType& InValue)
    {
        if (!Root)
        {
            Root = new RedBlackTreeNode(nullptr, nullptr, InValue, false);

            return;
        }
        Root = Add(Root, InValue);
        Root->bIsRed = false;
    }

    uint64_t Num() const
    {
        if (Root)
        {
            return Root->Size;
        }
        return 0;
    }

    void Clear()
    {
        if (Root)
        {
            delete Root;
            Root = nullptr;
        }
    }

    bool Contains(const ValueType& InValue) const
    {
        RedBlackTreeNode* CurrentNode = Root;

        while (CurrentNode)
        {
            if (InValue < CurrentNode->Value)
            {
                CurrentNode = CurrentNode->Left;
            }
            else if (CurrentNode->Value < InValue)
            {
                CurrentNode = CurrentNode->Right;
            }
            else
            {
                return true;
            }
        }

        return false;
    }

    void Remove(const ValueType& InValue)
    {
        if ((!IsRed(Root->Left)) && (!IsRed(Root->Right)))
        {
            Root->bIsRed = true;
        }

        Root = Remove(Root, InValue);

        if (Root)
        {
            Root->bIsRed = false;
        }
    }
protected:
    RedBlackTreeNode* Root = nullptr;


    bool IsRed(RedBlackTreeNode* Node)
    {
        if (Node && Node->bIsRed)
        {
            return true;
        }

        return false;
    }

    //return result parent node
    RedBlackTreeNode* Add(RedBlackTreeNode* InParentNode, const ValueType& InValue)
    {
        if (!InParentNode)
        {
            return  new RedBlackTreeNode(nullptr, nullptr, InValue, true);
        }

        if (InValue < InParentNode->Value)
        {
            InParentNode->Left = Add(InParentNode->Left, InValue);
        }
        else if(InParentNode->Value < InValue)
        {
            InParentNode->Right = Add(InParentNode->Right, InValue);
        }
        else
        {
            return InParentNode; //already exist donot add
        }

        InParentNode->Size++;

        RedBlackTreeNode* ResultPrent = InParentNode;

        //try rotate
        if (
            IsRed(ResultPrent->Right)
            && (!IsRed(ResultPrent->Left)))
        {
            ResultPrent = RotateLeft(ResultPrent);
        }

        if (IsRed(ResultPrent->Left)
            && IsRed(ResultPrent->Left->Left))
        {
            ResultPrent = RotateRight(ResultPrent);
        }

        if (IsRed(ResultPrent->Right)
            && IsRed(ResultPrent->Left))
        {
            FlipColor(ResultPrent);
        }

        return ResultPrent;

    }

    RedBlackTreeNode* Remove(RedBlackTreeNode* InParentNode, const ValueType& InValue)
    {

        if (InValue < InParentNode->Value)
        {
            if (!IsRed(InParentNode->Left) && !IsRed(InParentNode->Left->Left))
            {
                InParentNode = MoveRedLeft(InParentNode);
            }
            InParentNode->Left = Remove(InParentNode->Left, InValue);

        }
        else
        {
            if (IsRed(InParentNode->Left))
            {
                InParentNode = RotateRight(InParentNode);
            }


            if (!(InParentNode->Value < InValue)  && !InParentNode->Right)
            {
                return nullptr;
            }

            if (!IsRed(InParentNode->Right) && !IsRed(InParentNode->Right->Left))
            {
                InParentNode = MoveRedRight(InParentNode);
            }

            if (!(InParentNode->Value < InValue))
            {
                {
                    RedBlackTreeNode* MinNode = Min(InParentNode->Right);
                    InParentNode->Value = MinNode->Value;
                }
                InParentNode->Right = RemoveMin(InParentNode->Right);

            }
            else
            {
                InParentNode->Right = Remove(InParentNode->Right, InValue);
            }

        }

        return Balance(InParentNode);
    
    
    }

    RedBlackTreeNode* RemoveMin(RedBlackTreeNode* InParentNode)
    {
        if (!InParentNode->Left)
        {
            return nullptr;
        }

        if (!IsRed(InParentNode->Left) && !IsRed(InParentNode->Left->Left))
        {
            InParentNode = MoveRedLeft(InParentNode);
        }

        InParentNode->Left = RemoveMin(InParentNode->Left);

        return Balance(InParentNode);
    }

    RedBlackTreeNode* Min(RedBlackTreeNode* InParentNode)
    {
        if (InParentNode->Left)
        {
            return Min(InParentNode->Left);
        }

        return InParentNode;
    }

    //asume InParentNode is red
    //let left or left'child become red
    RedBlackTreeNode* MoveRedLeft(RedBlackTreeNode* InParentNode)
    {
        FlipColor(InParentNode);
        if (IsRed(InParentNode->Right->Left))
        {
            InParentNode->Right = RotateRight(InParentNode->Right);
            InParentNode = RotateLeft(InParentNode);
            FlipColor(InParentNode);
        }
        return InParentNode;
    }

    RedBlackTreeNode* MoveRedRight(RedBlackTreeNode* InParentNode)
    {
        FlipColor(InParentNode);
        if (IsRed(InParentNode->Left->Left))
        {
            InParentNode = RotateRight(InParentNode);
            FlipColor(InParentNode);
        }
        return InParentNode;
    }

    RedBlackTreeNode* Balance(RedBlackTreeNode* InParentNode)
    {
        RedBlackTreeNode* ResultPrent = InParentNode;

        if (IsRed(ResultPrent->Right))
        {
            ResultPrent = RotateLeft(ResultPrent);
        }

        if (
            IsRed(ResultPrent->Right)
            && (!IsRed(ResultPrent->Left)))
        {
            ResultPrent = RotateLeft(ResultPrent);
        }

        if (IsRed(ResultPrent->Left)
            && IsRed(ResultPrent->Left->Left))
        {
            ResultPrent = RotateRight(ResultPrent);
        }

        if (IsRed(ResultPrent->Right)
            && IsRed(ResultPrent->Left))
        {
            FlipColor(ResultPrent);
        }

        return ResultPrent;
    }

    //rotate right red node to Left
    RedBlackTreeNode* RotateLeft(RedBlackTreeNode* InNode)
    {
       // assert(InNode->Right->bIsRed);
        RedBlackTreeNode* Result = InNode->Right;

        InNode->Right = Result->Left;
        Result->Left = InNode;

        Result->bIsRed = InNode->bIsRed;
        InNode->bIsRed = true;

        Result->Size = InNode->Size;
        InNode->Size = (InNode->Left?InNode->Left->Size:0) + (InNode->Right ? InNode->Right->Size : 0) + 1;

        return Result;
    }

    //rotate red left node to right
    RedBlackTreeNode* RotateRight(RedBlackTreeNode* InNode)
    {
        //assert(InNode->Left->bIsRed);
        RedBlackTreeNode* Result = InNode->Left;

        InNode->Left = Result->Right;
        Result->Right = InNode;

        Result->bIsRed = InNode->bIsRed;
        InNode->bIsRed = true;

        Result->Size = InNode->Size;
        InNode->Size = (InNode->Left ? InNode->Left->Size : 0) + (InNode->Right ? InNode->Right->Size : 0) + 1;

        return Result;
    }

    //if both left and right are red,flip it
    void FlipColor(RedBlackTreeNode* InNode)
    {
        InNode->bIsRed = !InNode->bIsRed;
        InNode->Left->bIsRed = !InNode->Left->bIsRed;
        InNode->Right->bIsRed = !InNode->Right->bIsRed;
    }
};